#include <stdlib.h>

#include "log.h"
#include "memmap.h"
#include "hle/__SysMemUserForUser.h"

#include "elfreader.h"

#define ET_EXEC 2
#define ET_PRX  0xffa0

#define SHT_REL 9
#define SHT_PSPREL 0x700000a0

#define SHF_ALLOC 0x2

#define R_MIPS_NONE 0
#define R_MIPS_32 2
#define R_MIPS_26 4
#define R_MIPS_HI16 5
#define R_MIPS_LO16 6
#define R_MIPS_GPREL 7

void elfr_getMaxVaddr(ElfReader *reader)
{
    u32 i;
    u32 nseg = reader->header.e_phnum;
    for (i = 0; i < nseg; i++)
    {   
        Elf32_Phdr *p = &reader->segments[i];
        if (p->p_type == 1)
            if (p->p_vaddr + p->p_memsz > reader->maxVaddr)
                reader->maxVaddr = p->p_vaddr + p->p_memsz;
    }
}

void elfr_loadPrograms(ElfReader *reader)
{
    u32 i, j;
    u32 nseg = reader->header.e_phnum;
    BlkAlloc *part;
    if (blkAlloc_addrInBlock(&userMemory, reader->vaddr)
     || blkAlloc_addrInBlock(&userMemory, reader->header.e_entry))
        part = &userMemory;
    else if (blkAlloc_addrInBlock(&kernelMemory, reader->vaddr)
          || blkAlloc_addrInBlock(&kernelMemory, reader->header.e_entry))
        part = &kernelMemory;
    else {
        _log(ERR, ELF, "Can't copy segments to memory address 0x%08x!", reader->vaddr);
        return;
    }

    _log(INF, ELF, "%d segments:", nseg);
    /* loads segments into memory */
    for (i = 0; i < nseg; i++)
    {
        Elf32_Phdr *p = &reader->segments[i];

        _log(INF, ELF, "* Type: %x Vaddr: %08x Offset: %08x Filesz: %x Memsz: %x ", p->p_type, p->p_vaddr, p->p_offset, p->p_filesz, p->p_memsz);
        if (p->p_type == 1)
        {
            blkAlloc_allocAt(part, reader->vaddr + p->p_vaddr, p->p_memsz, "ELF", 0);
            for (j = 0; j < p->p_filesz; j++)
                mem_write8(p->p_vaddr + reader->vaddr + j, *(reader->base + p->p_offset + j));
            for (; j < p->p_memsz; j++)
                mem_write8(p->p_vaddr + reader->vaddr + j, 0);
            _log(INF, ELF, "Loadable Segment Copied to %08x, size %08x, end %08x, bss %08x",
                           p->p_vaddr + reader->vaddr, p->p_memsz, p->p_vaddr + reader->vaddr + p->p_memsz, p->p_memsz - p->p_filesz);
            if (p->p_vaddr + p->p_memsz > reader->maxVaddr)
                reader->maxVaddr = p->p_vaddr + p->p_memsz;
        }
    }
}

void elfr_loadSections(ElfReader *reader)
{
    /* loads sections */
    u32 nsect = reader->header.e_shnum;
    u32 i;
    _log(INF, ELF, "%d sections:", nsect);

    _log(INF, ELF, "______________________________________________________________________");
    _log(INF, ELF, "|%2s|%-35s|%-8s|%-8s|%-6s|%-4s|", "ID", "Name", "Address", "Size", "Offset", "Link");
    _log(INF, ELF, "|--------------------------------------------------------------------|");

    for (i = 0; i < nsect; i++) {
        Elf32_Shdr s = reader->sections[i];
        _log(INF, ELF, "|%02d|%-35s|%08x|%08x|%06x|%-4d|", i, elfr_getSectName(reader, i), s.sh_addr, s.sh_size, s.sh_offset, s.sh_link);
    }

    _log(INF, ELF, "+--------------------------------------------------------------------+");
}

void elfr_loadRelocation(ElfReader *reader, Elf32_Rel *rels, u32 numRelocs)
{
    u32 r, t;

    for (r = 0; r < numRelocs; r++)
    {
        u32 info = rels[r].r_info;
        u32 addr = rels[r].r_offset;

        u8 type = info & 0xf;

        u8 readwrite = (info >> 8) & 0xff;
        u8 relative  = (info >> 16) & 0xff;

        u32 op;
        u32 relocateTo;

        addr += reader->vaddr + reader->segments[readwrite].p_vaddr;

        op = mem_read32(addr);

        relocateTo = reader->vaddr + reader->segments[relative].p_vaddr;

        switch (type)
        {
        case R_MIPS_NONE:
            break;

        case R_MIPS_32:
            op += relocateTo;
            break;

        case R_MIPS_26:
            op = (op & 0xfc000000) | (((op & 0x03ffffff) + (relocateTo >> 2)) & 0x3fffffff);
            break;

        case R_MIPS_HI16:
            {
                u32 cur = (op & 0xffff) << 16;
                u16 hi = 0;
                s16 lo;

                for (t = r + 1; t < numRelocs; t++)
                {
                    if ((rels[t].r_info & 0xf) == R_MIPS_LO16)
                    {
                        u32 corrLoAddr = reader->vaddr + reader->segments[readwrite].p_vaddr + rels[t].r_offset;

                        lo = mem_read32(corrLoAddr) & 0xffff;
                        cur += lo;
                        cur += relocateTo;
                        lo = (cur & 0xffff);
                        hi = (cur - lo) >> 16;
                        break;
                    }
                }

                op = (op & 0xffff0000) | (hi);
            }
            break;

        case R_MIPS_LO16:
            op = (op & 0xffff0000) | ((op + relocateTo) & 0xffff);
            break;

        case R_MIPS_GPREL:
            _log(WRN, ELF, "GP reloc at %08x", addr);
            break;

        default:
            _log(ERR, ELF, "Unknown relocation %d at address %08x", type, addr);
            break;
        }

        mem_write32(addr, op);
    }
}

void elfr_loadRelocations(ElfReader *reader)
{
    u32 i;
    /* do relocations */
    _log(INF, ELF, "Relocations");

    for (i = 0; i < reader->header.e_phnum; i++)
    {
        Elf32_Phdr *p = &reader->segments[i];

        if (p->p_type == 0x700000a0)
        {
            u32 numRelocs = p->p_filesz / sizeof(Elf32_Rel);

            Elf32_Rel *rels = (Elf32_Rel*)(reader->base + p->p_offset);
            elfr_loadRelocation(reader, rels, numRelocs);
            return;
        }
        else if (p->p_type == 0x700000a1)
            _log(WRN, ELF, "Not implemented program relocation of type 0x%08x", p->p_type);
    }

    for (i = 0; i < reader->header.e_shnum; i++)
    {
        Elf32_Shdr *s = &reader->sections[i];

        if (s->sh_type == SHT_PSPREL)
        {
            u32 sectionToModify = s->sh_info;

            u32 numRelocs = s->sh_size / sizeof(Elf32_Rel);

            Elf32_Rel *rels = (Elf32_Rel*)(reader->base + s->sh_offset);
            if (!(reader->sections[sectionToModify].sh_flags & SHF_ALLOC)) {
                _log(WRN, ELF, "Trying to relocate non-loaded section %s", elfr_getSectName(reader, sectionToModify));
                continue;
            }

            elfr_loadRelocation(reader, rels, numRelocs);
        }
        else if (s->sh_type == SHT_REL)
            _log(WRN, ELF, "Traditional relocation section found in %s, not relocating", elfr_getSectName(reader, i));
    }
}

ElfReader elfr_init(void *ptr)
{
    ElfReader reader;
    reader.base = (char*)ptr;
    reader.header = *(Elf32_Ehdr*)ptr;
    reader.segments = (Elf32_Phdr*)(reader.base + reader.header.e_phoff);
    reader.sections = (Elf32_Shdr*)(reader.base + reader.header.e_shoff);
    reader.isPrx = ((reader.header.e_type & 0xffff) == ET_PRX);
    reader.maxVaddr = 0;
    elfr_getMaxVaddr(&reader);
    return reader;
}

void elfr_load(ElfReader *reader, u32 vaddr)
{
    reader->vaddr = reader->isPrx ? vaddr : 0;

    elfr_loadPrograms(reader);
    blkAlloc_listBlocks(&kernelMemory);

    elfr_loadSections(reader);

    if (reader->isPrx)
        elfr_loadRelocations(reader);

    _log(INF, ELF, "ELF file correctly loaded at %08x", vaddr);
}

const char *elfr_getSectName(ElfReader *reader, s32 section)
{
    return reader->base + reader->sections[reader->header.e_shstrndx].sh_offset + reader->sections[section].sh_name;
}

SectionID elfr_getSectByName(ElfReader *reader, const char *name)
{
    s32 i;
    for (i = 0; i < reader->header.e_shnum; i++)
        if (strcmp(name, elfr_getSectName(reader, i)) == 0)
            return i;
    return -1;
}

